/**************************************************************************
* Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
*                                                                        *
* Author: The ALICE Off-line Project.                                    *
* Contributors are mentioned in the code where appropriate.              *
*                                                                        *
* Permission to use, copy, modify and distribute this software and its   *
* documentation strictly for non-commercial purposes is hereby granted   *
* without fee, provided that the above copyright notice appears in all   *
* copies and that both the copyright notice and this permission notice   *
* appear in the supporting documentation. The authors make no claims     *
* about the suitability of this software for any purpose. It is          *
* provided "as is" without express or implied warranty.                  *
**************************************************************************/
//
//  Xianguo Lu <lu@physi.uni-heidelberg.de>
//

#include <TFile.h>
#include <TH2I.h>
#include <TH2D.h>
#include <TList.h>
#include <THnSparse.h>
#include <TSystem.h>

#include "AliAnalysisManager.h"
#include "AliESDfriend.h"
#include "AliESDtrackCuts.h"
#include "AliESDInputHandler.h"
//#include "AliTPCseed.h"
#include "AliTrackerBase.h"
#include "AliTRDseedV1.h"
#include "AliTRDtrackV1.h"
#include "AliCentrality.h"
#include "AliESDv0KineCuts.h"
#include "AliESDv0.h"

//#include "AliTRDdEdxBaseUtils.h"
//#include "AliTPCCosmicUtils.h"
#include "AliAnalysisTaskTRDdEdx.h"

#include "AliPIDResponse.h"

AliAnalysisTaskTRDdEdx::AliAnalysisTaskTRDdEdx()://const TString chunk):
  AliAnalysisTaskSE("TRDdEdx")
  //, fChunk(chunk)
  , fStream(0x0)
  , fESDEvent(0x0)
  , fESDpid(0x0)
  , fESDfriend(0x0)
  , fpid(0x0)
  , fList(0x0)
  , fhh(0x0)
{
  //
  //
  //

  DefineOutput(1, TTree::Class());
  DefineOutput(2, TList::Class());

  fList = new TList;
  fList->SetOwner();//------>>> Very important!!!!!!!
  
  fhh = new TH1I("hh","",20,0,20);
  fList->Add(fhh);

  //fStream = new TTreeSRedirector(Form("TaskTRDdEdxTree_%s.root", fChunk.Data()));
  fStream = new TTreeStream("tree");
  fTree=fStream->GetTree();
}

//____________________________________________________________
AliAnalysisTaskTRDdEdx::~AliAnalysisTaskTRDdEdx(){
  //
  // Destructor
  //
  //delete fStream;

  /* can not be deleted in chain, they are in input handler!
  delete fESDEvent;
  delete fESDfriend;
  delete fESDpid;
  */

  fList->Clear();
  delete fList;
}

void AliAnalysisTaskTRDdEdx::UserCreateOutputObjects()
{
  //
  //
  //

  PostData(1, fTree);
  PostData(2, fList);
}

void AliAnalysisTaskTRDdEdx::UserExec(Option_t *)
{
  //
  //Execute
  // 
  fhh->Fill(0);

  Int_t mag = -999;
  if(!CutESD(mag))
    return;

  fhh->Fill(1);

  if(!CutEvent())
    return;

  fhh->Fill(2);

  const Int_t ntrk0 = fESDEvent->GetNumberOfTracks();
  AliESDtrack * trks[ntrk0];
  for(Int_t ii=0; ii<ntrk0; ii++){
    trks[ii] = 0x0;
  }
  Int_t nsel = 0;
  CutTrack(nsel, trks);
  if(!nsel)
    return;

  fhh->Fill(3);

  Int_t pdgs[ntrk0];
  GetV0type(nsel, pdgs);

  Double_t pbcen = -999;
  AliCentrality *esdCentrality = fESDEvent->GetCentrality();
  if(esdCentrality){
    pbcen = esdCentrality->GetCentralityPercentile("V0M");
  }
  
  Int_t run = fESDEvent->GetRunNumber();
  for(Int_t itrk=0; itrk<nsel ; itrk++){
    AliESDtrack * esdtrack = trks[itrk];

    if(!esdtrack->GetTRDntracklets())
      continue;
    
    Int_t trdnch  = esdtrack->GetTRDNchamberdEdx();
    if(trdnch<1)
      continue;

    /*
    AliTRDtrackV1 * trdv1 = AliTRDdEdxBaseUtils::GetTRDtrackV1(esdtrack);
    if(!trdv1)
      continue;

    Int_t istack = -999, isector = -999;
    Double_t pTRD = -999;
    if(!AliTRDdEdxBaseUtils::GetFirstSectorStackMomentum(trdv1, isector, istack, pTRD)){
      continue;
    }
    */
    
    /*
    AliTPCseed * tpcseed = AliTPCCosmicUtils::GetTPCseed(esdtrack);
    if(!tpcseed)
      continue;
    */

    Double_t pidITS[5], pidTPC[5], pidTOF[5];
    esdtrack->GetITSpid(pidITS);
    esdtrack->GetTPCpid(pidTPC);
    esdtrack->GetTOFpid(pidTOF);
    
    Double_t tpcele    = fpid->NumberOfSigmasTPC( esdtrack, AliPID::kElectron );
    Double_t tpcpion   = fpid->NumberOfSigmasTPC( esdtrack, AliPID::kPion );
    Double_t tpckaon   = fpid->NumberOfSigmasTPC( esdtrack, AliPID::kKaon );
    Double_t tpcproton = fpid->NumberOfSigmasTPC( esdtrack, AliPID::kProton );
  
    //===>
    Double_t tofpion = -999;
    Double_t tofkaon = -999;
    Double_t tofproton = -999;
    Double_t tofele = -999;

    //http://alisoft.cern.ch/viewvc/trunk/STEER/STEERBase/AliPIDResponse.cxx?view=markup&root=AliRoot
    //check kTOFout, kTIME, mismatch
    const Int_t dummynSpe = 4;
    Double_t dummypro[4]={-999,-999,-999,-999};
    AliPIDResponse::EDetPidStatus tofstatus = fpid->ComputeTOFProbability(esdtrack, dummynSpe, dummypro);
    if( tofstatus == AliPIDResponse::kDetPidOk ){
      tofpion     = fpid->NumberOfSigmasTOF( esdtrack, AliPID::kPion);
      tofkaon     = fpid->NumberOfSigmasTOF( esdtrack, AliPID::kKaon);
      tofproton   = fpid->NumberOfSigmasTOF( esdtrack, AliPID::kProton);
      tofele = fpid->NumberOfSigmasTOF( esdtrack, AliPID::kElectron);
    }
    //<===

    /*
    const Double_t probCut = 1;
    if     ( esdtrack->TestBit(BIT(14)) && (pidTPC[AliPID::kElectron]+pidTOF[AliPID::kElectron])> probCut) pidV0 = AliPID::kElectron; 
    else if( esdtrack->TestBit(BIT(15)) && (pidTPC[AliPID::kPion]     +pidTOF[AliPID::kPion])   > probCut) pidV0 = AliPID::kPion; 
    else if( esdtrack->TestBit(BIT(16)) && (pidTPC[AliPID::kProton]  +pidTOF[AliPID::kProton])  > probCut) pidV0 = AliPID::kProton;
    */

    //in the tender it is the status bits of the tracks 14 - 16 for electrons, pions and protons
    Int_t tmpbit = 0;
    if(esdtrack->TestBit(BIT(14))) tmpbit+=1;
    if(esdtrack->TestBit(BIT(15))) tmpbit+=2;
    if(esdtrack->TestBit(BIT(16))) tmpbit+=4;

    Int_t pidV0 = -999;

    /*
    if(tmpbit == 1) pidV0 = AliPID::kElectron; 
    if(tmpbit == 2) pidV0 = AliPID::kPion; 
    if(tmpbit == 4) pidV0 = AliPID::kProton;
    */
    pidV0 = tmpbit;

    /*
    if(pidV0<-990){
      continue;
    }
    */

    Int_t kinV0 = pdgs[itrk];
    Int_t charge = esdtrack->Charge();
    Double_t p0 = esdtrack->P();
    Double_t eta = esdtrack->Eta();   

    Double_t pTPC = esdtrack->GetTPCmomentum();
    Double_t tpcsig  = esdtrack->GetTPCsignal();
    Int_t tpcncls = esdtrack->GetTPCsignalN();

    Double_t pTRD = -999;
    for(Int_t ich=0; ich<6; ich++){
      pTRD = esdtrack->GetTRDmomentum(ich);
      if(pTRD>0)
        break;
    }

    Double_t trdsig = esdtrack->GetTRDsignal();
    Int_t trdncls = esdtrack->GetTRDNclusterdEdx();


    Double_t sigmaTRD[5], deltaTRD[5], ratioTRD[5];
    AliPID::EParticleType types[]={AliPID::kElectron, AliPID::kMuon, AliPID::kPion, AliPID::kKaon, AliPID::kProton};
    for(Int_t itrdpid=0; itrdpid<5; itrdpid++){
      sigmaTRD[itrdpid]= fpid->NumberOfSigmas(AliPIDResponse::kTRD, esdtrack, types[itrdpid]);
      deltaTRD[itrdpid]= fpid->GetSignalDelta(AliPIDResponse::kTRD, esdtrack, types[itrdpid]);
      ratioTRD[itrdpid]= fpid->GetSignalDelta(AliPIDResponse::kTRD, esdtrack, types[itrdpid], kTRUE);
    }

    (*fStream)<<"tree"<<
      "run="<<run<<
      "mag="<<mag<<
      "charge="<<charge<<
      "pbcen="<<pbcen<<

      "p0="<<p0<<
      "eta="<<eta<<

      /*
      "esdtrack="<<esdtrack<<
      "trdv1="<<trdv1<<
      "tpcseed="<<tpcseed<<
      */

      "tpcele="<<tpcele<<
      "tpcpion="<<tpcpion<<
      "tpckaon="<<tpckaon<<
      "tpcproton="<<tpcproton<<
   
      "tofele="<<tofele<<
      "tofpion="<<tofpion<<
      "tofkaon="<<tofkaon<<
      "tofproton="<<tofproton<<

      "pTPC="<<pTPC<<
      "tpcsig="<<tpcsig<<
      "tpcncls="<<tpcncls<<

      "pTRD="<<pTRD<<
      "trdsig="<<trdsig<<
      "trdnch="<<trdnch<<
      "trdncls="<<trdncls<<

      "kinV0="<<kinV0<<
      "pidV0="<<pidV0;

    for(Int_t ipid=0; ipid<5; ipid++){
      (*fStream)<<"tree"<<
        Form("pidITS%d=",ipid)<<pidITS[ipid]<<
        Form("pidTPC%d=",ipid)<<pidTPC[ipid]<<
        Form("pidTOF%d=",ipid)<<pidTOF[ipid]<<
        Form("sigmaTRD%d=",ipid)<<sigmaTRD[ipid]<<
        Form("deltaTRD%d=",ipid)<<deltaTRD[ipid]<<
        Form("ratioTRD%d=",ipid)<<ratioTRD[ipid];
    }
      
    (*fStream)<<"tree"<<
      "\n";

  } 

  fhh->Fill(4);
}

void AliAnalysisTaskTRDdEdx::Terminate(Option_t *) 
{
  //
  //do when Terminate
  //
}

//========================================================================================
//========================================================================================

Int_t AliAnalysisTaskTRDdEdx::GetPID(const AliESDtrack *trk)
{
  //
  //in the tender it is the status bits of the tracks 14 - 16 for electrons, pions and protons
  //
  Int_t ipar = -999;

  const Double_t nsigma = 3;
  const Double_t tpcele    = fESDpid->NumberOfSigmasTPC( trk, AliPID::kElectron );
  const Double_t tpcpion   = fESDpid->NumberOfSigmasTPC( trk, AliPID::kPion );
  const Double_t tpckaon   = fESDpid->NumberOfSigmasTPC( trk, AliPID::kKaon );
  const Double_t tpcproton = fESDpid->NumberOfSigmasTPC( trk, AliPID::kProton );

  const Bool_t ka = !(trk->GetStatus() & AliESDtrack::kTOFmismatch);
  const Bool_t kb =  (trk->GetStatus() & AliESDtrack::kTOFpid);
  const Bool_t ktof = ka && kb;
  if(ktof){
    const Double_t tofelectron = fESDpid->NumberOfSigmasTOF( trk, AliPID::kElectron, fESDpid->GetTOFResponse().GetTimeZero() );
    const Double_t tofpion     = fESDpid->NumberOfSigmasTOF( trk, AliPID::kPion,     fESDpid->GetTOFResponse().GetTimeZero() );
    const Double_t tofkaon     = fESDpid->NumberOfSigmasTOF( trk, AliPID::kKaon,     fESDpid->GetTOFResponse().GetTimeZero() );
    const Double_t tofproton   = fESDpid->NumberOfSigmasTOF( trk, AliPID::kProton,   fESDpid->GetTOFResponse().GetTimeZero() );

    if((fabs(tpcele)   <nsigma)  && (fabs(tofelectron)<nsigma) && trk->TestBit(BIT(14)) )
      ipar = AliPID::kElectron;

    if((fabs(tpcpion)  <nsigma)  && (fabs(tofpion)    <nsigma) && trk->TestBit(BIT(15)) )
      ipar = AliPID::kPion;

    if((fabs(tpckaon)  <1)  && (fabs(tofkaon)    <1) )
      ipar = AliPID::kKaon; 

    if((fabs(tpcproton)<nsigma)  && (fabs(tofproton)  <nsigma) && trk->TestBit(BIT(16)) )
      ipar = AliPID::kProton;
  }

  return ipar;
}

Bool_t AliAnalysisTaskTRDdEdx::CutEvent()
{
  Bool_t kpr0 = kTRUE;
  const AliESDVertex *vertex = fESDEvent->GetPrimaryVertexTracks();
  if(vertex->GetNContributors()<1) {
    // SPD vertex
    vertex = fESDEvent->GetPrimaryVertexSPD();
    if(vertex->GetNContributors()<1) {
      // NO GOOD VERTEX, SKIP EVENT 
      kpr0 = kFALSE;
    }
  }
  const Bool_t kpriv = kpr0 && ( fabs(fESDEvent->GetPrimaryVertex()->GetZ())<10 ); 

  return kpriv;

}

void AliAnalysisTaskTRDdEdx::CutTrack(Int_t &nsel,  AliESDtrack *outtrk[])
{
  //
  //CutTrack
  //
  const Double_t etacut = 0.9;

  AliESDtrackCuts esdTrackCuts;//i.e. GetStandardITSTPCTrackCuts2010(kTRUE);

  //--->
  // TPC  
  esdTrackCuts.SetMinNCrossedRowsTPC(70);
  //esdTrackCuts.SetMinRatioCrossedRowsOverFindableClustersTPC(0.8);

  esdTrackCuts.SetMaxChi2PerClusterTPC(4);
  esdTrackCuts.SetAcceptKinkDaughters(kFALSE);
  esdTrackCuts.SetRequireTPCRefit(kTRUE);
  // ITS
  esdTrackCuts.SetRequireITSRefit(kTRUE);
  /*
  esdTrackCuts.SetClusterRequirementITS(AliESDtrackCuts::kSPD,
                                        AliESDtrackCuts::kAny);
  // 7*(0.0026+0.0050/pt^1.01)
  esdTrackCuts.SetMaxDCAToVertexXYPtDep("0.0182+0.0350/pt^1.01");

  esdTrackCuts.SetMaxDCAToVertexZ(2);
  esdTrackCuts.SetDCAToVertex2D(kFALSE);
  esdTrackCuts.SetRequireSigmaToVertex(kFALSE);
  */
  //---<
  
  //---
  esdTrackCuts.SetEtaRange(-etacut, etacut);
  //---------------------------------------------------------------<

  const TObjArray* seltrk = esdTrackCuts.GetAcceptedTracks(fESDEvent);
  nsel = seltrk->GetEntries();
  for(Int_t ii=0; ii<nsel; ii++){
    AliESDtrack *esdtrack = (AliESDtrack *)(seltrk->At(ii));
    outtrk[ii]=esdtrack;
  }
  delete seltrk;
}


Bool_t AliAnalysisTaskTRDdEdx::CutESD(Int_t & mag)
{
  //
  //general protection
  //
  fESDEvent = 0x0;
  fESDfriend = 0x0;

  AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler *>(AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
  if(esdH){  
    fESDEvent = esdH->GetEvent();
    fpid = esdH->GetPIDResponse();
  }

  if(!fpid){
    printf("AliAnalysisTaskpkRaa::CutInput error No ESD ESDpid\n"); exit(1);
  }

  // if(!fESDpid){
  //   printf("xlulog no ESDpid!!\n");
  //   return 0;
  // }

  if(!fESDEvent){
    printf("xlulog No ESD Event\n");
    exit(1);
  }
  
  esdH->SetActiveBranches("ESDfriend*");
  fESDfriend = (AliESDfriend *)fESDEvent->FindListObject("AliESDfriend");

  /*
  if(!fESDfriend){
    printf("xlulog No ESD friend\n");
    exit(1);
  }
  */
  if(fESDfriend){
    fESDEvent->SetESDfriend(fESDfriend);
  }
  //---------------------

  //in unit of kg
  const Double_t dm = fESDEvent->GetMagneticField();
  if( fabs(fabs(dm)- 5) > 0.1 && fabs(fabs(dm)- 1) > 0.1){
    printf("xlulog strange Bfield! %f\n", dm);
    exit(1);
  }

  if(dm>0)
    mag = (int) (dm+0.5);
  else
    mag = (int) (dm-0.5);

  return kTRUE;
}

// svn cat AliAnalysisTaskpkRaa.cxx -r 761
Int_t AliAnalysisTaskTRDdEdx::GetPIDbit(const Int_t rpdg)
{
  switch(TMath::Abs(rpdg)){
  case kElectron:
    return 1;
  case kPiPlus:
    return 2;
  case kProton:
    return 4;
  default:
    return 8;
  }
}

void AliAnalysisTaskTRDdEdx::GetV0type(const Int_t ntrk, Int_t pdgs[])
{
  //ntrk all tracks

  for(Int_t ii=0; ii<ntrk; ii++){
    pdgs[ii]= 0;
  }

  AliESDv0KineCuts v0cuts;
  //setevent will set primaryvertex automatically  
  v0cuts.SetEvent(fESDEvent);

  //for pp, cen all = -1
  v0cuts.SetMode(AliESDv0KineCuts::kPurity, AliESDv0KineCuts::kPP);

  for(Int_t iv0=0; iv0<fESDEvent->GetNumberOfV0s(); iv0++){
    AliESDv0 * esdv0 = fESDEvent->GetV0(iv0);
    if(!esdv0->GetOnFlyStatus())
      continue;

    Int_t pdgV0 = -999, pdgP = -999, pdgN = -999;
    if(v0cuts.ProcessV0(esdv0, pdgV0, pdgP, pdgN)){

      const Int_t itrkP = esdv0->GetPindex();
      const Int_t itrkN = esdv0->GetNindex();

      //if reused, pdgs[]==fkBad
      pdgs[itrkP]+=GetPIDbit(pdgP);
      pdgs[itrkN]+=GetPIDbit(pdgN);
    }
  }
}
